#!/usr/bin/python

"""
This program draws MUMMER alignment results and shows connecting segments
based on % identity.

Usage: dissertation_DrawMUMMER.py g1.gbk g2.gbk g1.ptt g2.ptt cogs.t.list mummer.coord 
Examples:
Go to this directory:
/host/Users/JS/UH-work/gloeobacter/final_work/comparisons

dissertation_DrawMUMMERwithPttZoomRegion.py NC_007776.gbk NC_007775.gbk NC_007776.ptt NC_007775.ptt orthologs/orthomcl/cogs.t.list NC_007776_vs_NC_007775.coords 0 30000
dissertation_DrawMUMMERwithPttZoomRegion.py NC_007516.gbk NC_007513.gbk NC_007516.ptt NC_007513.ptt orthologs/orthomcl/cogs.t.list NC_007516_vs_NC_007513.coords 0 30000
dissertation_DrawMUMMERwithPttZoomRegion.py NC_011748.gbk NC_008253.gbk NC_011748.ptt NC_008253.ptt orthologs/orthomcl/cogs.t.list NC_011748_vs_NC_008253.coords 0 30000
dissertation_DrawMUMMERwithPttZoomRegion.py NC_002737.gbk NC_007297.gbk NC_002737.ptt NC_007297.ptt orthologs/orthomcl/cogs.t.list NC_002737_vs_NC_007297.coords 0 30000
dissertation_DrawMUMMERwithPttZoomRegion.py NC_007516.gbk NC_007513.gbk NC_007516.ptt NC_007513.ptt orthologs/orthomcl/cogs.t.list NC_007516_vs_NC_007513.coords 0 30000
Steps:
1. Download fna, gbk, and ptt files
2. Run MUMMER
3. Run this program

Author: Jimmy Saw
Date of last update: 04-25-2012
NOTE: NOT WORKING AS INTENDED. NEED TO WORK ON MORE.

"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC
import matplotlib.patches as mpatch

#Regex and other stuffs
cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')
#pttcog = re.compile('(COG\d+)(\w).*')
pttcog = re.compile('(COG\d{4})(\w).*')

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#D2B9D3',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
}



##Genome 1 Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    if feat.type == 'CDS':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 2 Genbank file
g2seq = SeqIO.read(sys.argv[2], "gb")
g2length = len(g2seq.seq)
g2featdict = {}
for feat in g2seq.features:
    if feat.type == 'CDS':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 1 ptt file
g1cogs = {}
g1pttfile = open(sys.argv[3], "rU")
g1ptt = g1pttfile.readlines()
for line in g1ptt[3:]:
    c = line.split('\t')
    g1ltag = c[5]
    g1gene = c[4]
    #g1cogcat = '-'
    g1cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        #g1cogcat = p.group(2)
        g1cog = p.group(1)
    #g1cogs[g1ltag] = g1cogcat
    g1cogs[g1ltag] = g1cog

##Genome 2 ptt file
g2cogs = {}
g2pttfile = open(sys.argv[4], "rU")
g2ptt = g2pttfile.readlines()
for line in g2ptt[3:]:
    c = line.split('\t')
    g2ltag = c[5]
    g2gene = c[4]
    #g2cogcat = '-'
    g2cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        #g2cogcat = p.group(2)
        g2cog = p.group(1)
    #g2cogs[g2ltag] = g2cogcat
    g2cogs[g2ltag] = g2cog

cogcatfile = open(sys.argv[5], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

cogcatfile.close()

spanx1 = int(sys.argv[7])
spanx2 = int(sys.argv[8])

glist = []

genome1x = [spanx1, spanx2]
genome1y = [2, 2]
genome2x = [spanx1, spanx2]
genome2y = [12, 12]

largergenome = 0

if g1length > g2length:
    largergenome = g1length
else:
    largergenome = g2length

##Start plotting
fig = plt.figure(1, figsize=(14,5))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)

ax1.axis([0, largergenome, 0, 14])
#ax1.axis([spanx1, spanx2, 0, 14])

for k, v in g1featdict.iteritems():
    g1feat = v
    g1locustag = g1feat.qualifiers['locus_tag'][0]
    g1start = g1feat.location._start.position
    g1stop = g1feat.location._end.position
    g1size = g1stop - g1start + 1
    g1mid = g1start + ((g1stop - g1start) / 2.0)
    g1desc = g1feat.qualifiers['product'][0]
    g1gene = ""
    if g1feat.qualifiers.has_key('gene'):
        g1gene = g1feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g1gene = g1feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = '#D3D3D3' #base color
    if g1locustag in g1cogs:
        if g1cogs[g1locustag] != '-':
            #cogcolor = cogdict[g1cogs[g1locustag]]
            cogcolor = cogdict[cogcatdict[g1cogs[g1locustag]]]
    if g1feat.type == 'tRNA':
        cogcolor = '#800000'
    if g1feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g1feat.strand == -1:
        if spanx1 <= g1start <= spanx2 and spanx1 <= g1stop <= spanx2:
            #print "g1start: ", g1start
            rect = Rectangle((g1start, 4.0), g1size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
            plt.gca().add_patch(rect)
            ax1.text(g1mid, 4.5, g1gene, fontsize=8, color='black', rotation=45)
    else:
        if spanx1 <= g1start <= spanx2 and spanx1 <= g1stop <= spanx2:
            #print "g1start: ", g1start
            rect = Rectangle((g1start, 4.5), g1size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
            plt.gca().add_patch(rect)
            ax1.text(g1mid, 5.0, g1gene, fontsize=8, color='black', rotation=45)

#ax1.text(len(g1seq.seq)+10000, 4.2, "-", fontsize=10, color='black')
#ax1.text(len(g1seq.seq)+10000, 4.6, "+", fontsize=10, color='black')
#ax1.text(len(g2seq.seq)+10000, 9.2, "-", fontsize=10, color='black')
#ax1.text(len(g2seq.seq)+10000, 9.6, "+", fontsize=10, color='black')
ax1.annotate(g1seq.annotations['organism'], xy=(0.5, 0.05), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)
ax1.annotate(g2seq.annotations['organism'], xy=(0.5, 0.95), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)

##Parse MUMMER alignment file

nspanx1 = 0
nspanx2 = 0
g1starts = []
g1stops = []
g2starts = []
g2stops = []

"""
xdiff = 0

if g1spanx1 > g2spanx1:
    xdiff = g1spanx1 - g2spanx1
    padding = xdiff - g1starts[0]
else:
    xdiff = g2spanx1 - g1spanx1
    padding = g2starts[0] - xdiff
"""

mummerfile = open(sys.argv[6], "rU")
mfl = mummerfile.readlines()
for line in mfl[4:]:
    c = line.split('\t')
    g1start = int(c[0])
    g1stop = int(c[1])
    g2start = int(c[2])
    g2stop = int(c[3])
    ident = float(c[6])
    if spanx1 <= g1start and g1stop <= spanx2:
        g1starts.append(g1start)
        g1stops.append(g1stop)
        g2starts.append(g2start)
        g2stops.append(g2stop)
        #ax1.fill(x, y, color=fillcolor, alpha=0.2)

g1starts.sort()
g2stops.sort()
g2starts.sort()
g2stops.sort()

sdiff = 0
padding = 0

if len(g1starts) > 1:
    if g1starts[0] > g2starts[0]:
        padding = g1starts[0] - g2starts[0] #distance needed to pad to g2 minimum to flush coordinates
        sdiff = g1starts[0] - spanx1
    else:
        padding = g2starts[0] - g1starts[0]
        sdiff = spanx1 - g1starts[0]
nspanx1 = g2starts[0]
nspanx2 = g2stops[-1]


#print "spanx1, spanx2: ", spanx1, spanx2
#print "nspanx1, nspanx2: ", nspanx1, nspanx2
#print "padding: ", padding

for line in mfl[4:]:
    c = line.split('\t')
    g1start = int(c[0])
    g1stop = int(c[1])
    g1span = g1stop - g1start
    g2start = int(c[2])
    g2stop = int(c[3])
    g2span = g2stop - g2start
    ident = float(c[6])
    #reset genome 2 coordinates
    newg2start = 0
    newg2stop = 0
    if g2start >= g1start:
        #startdiff = g2start - spanx1
        #subtract = g2start - startdiff - spanx1 + sdiff
        #newg2start = g2start - subtract
        #newg2stop = g2stop - subtract
        diff1 = g2start - g1start
        diff2 = g1starts[0] - spanx1
        newg2start = g2start - diff1 + diff2
        newg2stop = newg2start + g2span
        print "g2start, newg2start, g2stop, newg2stop", g2start, newg2start, g2stop, newg2stop
    else:
        #startdiff = g1start - spanx1
        #subtract = g2start - startdiff - spanx1 + sdiff
        #newg2start = g2start + subtract
        #newg2stop = g2stop + subtract
        diff1 = g1start - g2start
        diff2 = g2starts[0] - spanx1
        newg2start = g2start + diff1 - diff2
        newg2stop = newg2start + g2span
        print "g2start, newg2start, g2stop, newg2stop", g2start, newg2start, g2stop, newg2stop
        
    fillcolor = '#AAAAAA'
    if ident >= 90:
        fillcolor = '#FF0000'
    elif ident >= 80:
        fillcolor = '#71C671'
    elif ident >= 70:
        fillcolor = '#7171C6'
    elif ident >= 60:
        fillcolor = '#CDB5CD'
    else:
        fillcolor = '#C5C1AA'
    #x = [g1start, g2start, g2stop, g1stop]
    x = [g1start, newg2start, newg2stop, g1stop]
    y = [5, 9, 9, 5]
    if spanx1 <= g1start and g1stop <= spanx2:
        #g1starts.append(g1start)
        #g1stops.append(g1stop)
        #g2starts.append(g2start)
        #g2stops.append(g2stop)
        ax1.fill(x, y, color=fillcolor, alpha=0.2)

#Parse genome 2 after checking MUMMER file
for k, v in g2featdict.iteritems():
    g2feat = v
    g2locustag = g2feat.qualifiers['locus_tag'][0]
    g2start = g2feat.location._start.position
    g2stop = g2feat.location._end.position
    g2size = g2stop - g2start + 1
    g2mid = g2start + ((g2stop - g2start) / 2.0)
    g2desc = g2feat.qualifiers['product'][0]
    g2gene = ""
    if g2feat.qualifiers.has_key('gene'):
        g2gene = g2feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g2gene = g2feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = '#D3D3D3' #base color
    if g2locustag in g2cogs:
        if g2cogs[g2locustag] != '-':
            #cogcolor = cogdict[g2cogs[g2locustag]]
            cogcolor = cogdict[cogcatdict[g2cogs[g2locustag]]]
    if g2feat.type == 'tRNA':
        cogcolor = '#800000'
    if g2feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g2gene = g2desc
    if g2feat.strand == -1:
        #if nspanx1 <= g2start <= nspanx2 and nspanx1 <= g2stop <= nspanx2:
        if g2start >= spanx1 and g2stop <= spanx2:
            newg2start = g2start + padding
            #print "newg2start, g2size: ", newg2start, g2size
            #print "g2start: ", g2start
            rect = Rectangle((newg2start, 9.0), g2size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
            plt.gca().add_patch(rect)
            ax1.text(g2mid, 9.5, g2gene, fontsize=8, color='black', rotation=45)
    else:
        #if nspanx1 <= g2start <= nspanx2 and nspanx1 <= g2stop <= nspanx2:
        if g2start >= spanx1 and g2stop <= spanx2:
            newg2start = g2start + padding
            #print "newg2start, g2size: ", newg2start, g2size
            #print "g2start: ", g2start
            rect = Rectangle((newg2start, 9.0), g2size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
            plt.gca().add_patch(rect)
            ax1.text(g2mid, 9.5, g2gene, fontsize=8, color='black', rotation=45)

"""
#This segment is an example of how to draw polygon connecting genomic regions
a = [2015000, 2018000, 2030000, 2020000]
a1 = [2015000, 2030000, 2018000, 2020000]
b = [6, 7, 7, 6]
ax1.fill(a1, b, 'r', alpha=0.2)
"""


frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)

plt.show()
